# 步骤1：获取数据（从数据库获取所有小说信息）
import asyncio
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))

from tortoise import Tortoise
from models.m import Novel, NovelStatus
from tool.sessing import TORTOISE_ORM


async def _get_novels_data_async() -> list:
    """异步获取所有小说数据，每本小说用list存储"""
    # 检查数据库是否已初始化，如果没有则初始化
    try:
        Tortoise.get_connection("default")
    except:
        await Tortoise.init(config=TORTOISE_ORM)

    try:
        # 查询所有小说，包含分类信息
        novels = await Novel.all().prefetch_related('category').all()

        novels_data = []

        for novel in novels:
            # 获取分类名称
            category_name = novel.category.name if novel.category else "未分类"

            # 获取状态文本
            status_text = {
                NovelStatus.SERIALIZING: "连载中",
                NovelStatus.COMPLETED: "已完结",
                NovelStatus.SUSPENDED: "暂停"
            }.get(novel.status, "未知")

            # 获取细节分类（tags字段，如果为空则返回空列表）
            detail_tags = novel.tags if novel.tags else []

            # 构建每本小说的数据列表（不再查询和存储章节标题）
            novel_info = [
                novel.title,  # 小说标题
                novel.author,  # 作者
                category_name,  # 分类
                status_text,  # 状态
                novel.word_count,  # 字数
                novel.chapter_count,  # 章节数
                float(novel.rating),  # 评分（转换为float）
                novel.view_count,  # 人气（浏览次数）
                novel.collect_count,  # 收藏
                detail_tags,  # 细节分类（标签列表）
                novel.description or ""  # 简介（如果为空则返回空字符串）
                # 移除了 chapter_titles（每章标题列表）
            ]

            novels_data.append(novel_info)

        return novels_data

    except Exception as e:
        print(f"获取小说数据时出错: {e}")
        return []


def get_data(file_path: str = None) -> list:
    """
    获取所有小说数据
    返回格式：[[小说1的信息列表], [小说2的信息列表], ...]
    每本小说信息列表包含：[标题, 作者, 分类, 状态, 字数, 章节数, 评分, 人气, 收藏, 细节分类, 简介]
    """
    # 使用 asyncio 运行异步函数
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    return loop.run_until_complete(_get_novels_data_async())


# 步骤2：分割文本（固定简单参数）
def split_text_grouped(novels_data: list, chunk_size: int = 600, chunk_overlap: int = 80) -> list:
    """
    将数据库中的小说数据分组，每本小说作为一组
    :param novels_data: 小说数据列表
    :param chunk_size: 每块大小
    :param chunk_overlap: 重叠大小
    :return: 分组后的数据
    """
    grouped_data = []
    
    for novel_info in novels_data:
        # 小说元信息
        novel_meta = {
            "title": novel_info[0],
            "author": novel_info[1],
            "category": novel_info[2],
            "status": novel_info[3],
            "word_count": novel_info[4],
            "chapter_count": novel_info[5],
            "rating": novel_info[6],
            "view_count": novel_info[7],
            "collect_count": novel_info[8],
            "tags": novel_info[9],
            "description": novel_info[10]
        }
        
        # 构建分组数据
        group = {
            "meta": novel_meta,
            "content": novel_info[10]  # 简介作为内容
        }
        
        grouped_data.append(group)
        
    return grouped_data


import pdfplumber
import re


def get_data2(pdf_path):
    """读取 PDF 文件并返回清洗清洗后的后的文本内容"""
    if not pdf_path.endswith('.pdf'):
        raise ValueError("参数必须是PDF文件路径（.pdf结尾）")

    raw_text = ""
    try:
        # 读取PDF所有页面文本
        with pdfplumber.open(pdf_path) as pdf:
            for page in pdf.pages:
                # 提取当前页文本，空页跳过
                page_text = page.extract_text()
                if page_text:
                    raw_text += page_text + "\n"  # 页间用换行分隔

        # 文本清洗：针对小说创作知识库的格式优化
        # 1. 去除多余空行（连续换行保留1个）
        cleaned = re.sub(r'\n+', '\n', raw_text).strip()

        # 2. 修复因PDF排版导致的单词/短语断裂（如"修仙\n体系" → "修仙体系"）
        # 匹配"换行+小写字母/数字"的情况，合并为一行
        cleaned = re.sub(r'(\w+)\n(\w+)', r'\1\2', cleaned)

        # 3. 去除PDF提取的冗余页脚/页码（如"Page 1/10"等）
        cleaned = re.sub(r'Page \d+/\d+', '', cleaned)

        # 4. 统一标题符号（确保##、###等标记连续，去除中间空格）
        cleaned = re.sub(r'#\s+#', '##', cleaned)
        cleaned = re.sub(r'##\s+#', '###', cleaned)

        return cleaned

    except FileNotFoundError:
        raise FileNotFoundError(f"未找到PDF文件：{pdf_path}")
    except Exception as e:
        raise RuntimeError(f"PDF读取失败：{str(e)}")





        
        
#用于过滤数据
def sjyhua(raw_data, raw_data2):
    """
    过滤数据函数
    :param raw_data: 来自数据库的小说数据列表
    :param raw_data2: 来自PDF的文本数据
    :return: filtered_data, filter_report (过滤后的数据, 过滤报告)
    """
    # 对于数据库数据，我们暂时不过滤，直接返回
    filtered_data = raw_data
    
    # 对于PDF数据，我们可以做一些基本的过滤
    # 例如，移除空行和多余的空白字符
    if isinstance(raw_data2, str):
        # 过滤PDF数据，移除过多的空行和空白字符
        lines = raw_data2.split('\n')
        filtered_lines = [line.strip() for line in lines if line.strip()]
        filtered_pdf_data = '\n'.join(filtered_lines)
        filter_report = filtered_pdf_data
    else:
        filter_report = raw_data2
    
    return filtered_data, filter_report


#语义分割
def split_text2(raw_data2):
    """
    对PDF文本进行语义分割
    :param raw_data2: PDF文本数据
    :return: 分割后的文本块列表
    """
    if not raw_data2:
        return []
    
    # 按章节标题分割文本
    # 查找章节标题模式（例如：## 一、小说类型与核心设定体系）
    sections = re.split(r'(#+\s*.*?[^\n]*)\n', raw_data2)
    
    # 清理和重组章节
    chunks = []
    i = 0
    while i < len(sections):
        if re.match(r'#+\s*.*', sections[i]):
            # 这是一个标题
            title = sections[i].strip()
            if i + 1 < len(sections):
                # 合并标题和内容
                content = sections[i + 1].strip() if sections[i + 1] else ""
                chunks.append(f"{title}\n{content}")
                i += 2
            else:
                chunks.append(title)
                i += 1
        else:
            # 独立段落
            if sections[i].strip():
                chunks.append(sections[i].strip())
            i += 1
    
    # 过滤掉空内容
    chunks = [chunk for chunk in chunks if chunk.strip()]
    
    return chunks


def vectorize(texts: list, grouped: list = None, vectorization_mode: str = "semantic") -> list:
    """
    使用阿里云百炼text-embedding-v4模型对文本进行向量化

    :param texts: 文本列表（来自PDF的分割文本块）
    :param grouped: 分组数据（来自数据库的小说信息）
    :param vectorization_mode: 向量量化模式
    :return: 向量列表
    """
    import os
    from openai import OpenAI

    # 从环境变量获取API密钥
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("未设置环境变量 DASHSCOPE_API_KEY")

    # 初始化OpenAI兼容客户端（连接阿里云百炼服务）
    client = OpenAI(
        api_key=api_key,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 存储所有文本和对应的向量
    all_texts = []
    all_vectors = []

    # 处理来自数据库的小说数据
    if grouped:
        for item in grouped:
            meta = item.get("meta", {})
            # 构建更丰富的描述文本，包含标题、作者和简介
            text_parts = []
            if meta.get("title"):
                text_parts.append(f"小说标题: {meta.get('title')}")
            if meta.get("author"):
                text_parts.append(f"作者: {meta.get('author')}")
            if meta.get("category"):
                text_parts.append(f"分类: {meta.get('category')}")
            if meta.get("status"):
                text_parts.append(f"状态: {meta.get('status')}")
            if meta.get("description"):
                text_parts.append(f"简介: {meta.get('description')}")

            # 添加标签信息
            if meta.get("tags"):
                tags_str = ", ".join(meta.get("tags", []))
                text_parts.append(f"标签: {tags_str}")

            text = " ".join(text_parts)
            all_texts.append(text)

    # 处理来自PDF的文本数据
    if texts:
        all_texts.extend(texts)

    # 对所有文本进行向量化（使用text-embedding-v4模型，支持批量处理）
    if not all_texts:
        return []

    # 阿里云text-embedding-v4模型限制：单次批处理最大10条文本
    batch_size = 10
    for i in range(0, len(all_texts), batch_size):
        batch = all_texts[i:i + batch_size]
        try:
            # 调用嵌入模型生成向量
            completion = client.embeddings.create(
                model="text-embedding-v4",  # 使用text-embedding-v4模型
                input=batch,
                dimensions=1024,  # 改回1024维以匹配ChromaDB期望的维度
                encoding_format="float"
            )

            if completion:
                # 提取当前批次的嵌入向量
                batch_vectors = [item.embedding for item in completion.data]
                all_vectors.extend(batch_vectors)
            else:
                print(f"批次 {i // batch_size + 1} 向量化失败，无返回结果")
                # 补充零向量作为占位符
                all_vectors.extend([[0.0] * 1024 for _ in range(len(batch))])

        except Exception as e:
            print(f"批次 {i // batch_size + 1} 向量化过程中发生异常: {e}")
            # 补充零向量作为占位符
            all_vectors.extend([[0.0] * 1024 for _ in range(len(batch))])

    return all_vectors